package com.yangzhe.algorithm.c035;

import java.util.PriorityQueue;

// 快速获得数据流的中位数的结构
public class Code05_MedianFinder_LeetCode {
    /**
     * 测试链接 : https://leetcode.cn/problems/find-median-from-data-stream/
     * 将数组分为两段，前半段用大根堆存储，后半段用小根堆存储
     * 数组个数为单数时，大根堆比小根堆多存储一个数字，中位数就是大根堆的堆顶
     * 数组个数为偶数时，中位数就是大根堆和小根堆堆顶再除以二
     */
    class MedianFinder {

        private PriorityQueue<Integer> minHeap;
        private PriorityQueue<Integer> maxHeap;

        public MedianFinder() {
            // PriorityQueue是小的数优先级高，也就是靠近堆顶，内部结构就是小根堆
            // 所以自然比较器 a - b，a = 2 , b = 5时， 2 - 5 = -3，程序认为a比b小，也就是'2'比'5'小，2会更靠近堆顶，形成小根堆
            minHeap = new PriorityQueue<>((a, b) -> (a - b));
            // 如果颠倒 b - a，a = 2, b= 5时， 5 - 2 = 3, 程序会认为 a比b大，也就是'2'比'5'大，5会更靠近堆顶，形成大根堆
            maxHeap = new PriorityQueue<>((a, b) -> (b - a));
        }

        public void addNum(int num) {
            // 比大根堆堆顶还小，也就是比前半段最后一个数字还小的，放在前半段，其余的放在后半段
            if (maxHeap.isEmpty() || num < maxHeap.peek()) {
                maxHeap.add(num);
            } else {
                minHeap.add(num);
            }

            // 如果后半段比前半段多
            if (!minHeap.isEmpty() && minHeap.size() > maxHeap.size()) {
                maxHeap.add(minHeap.poll());
            }

            // 如果前半段比后半段多出的数字超过1
            if (maxHeap.size() - minHeap.size() > 1) {
                minHeap.add((maxHeap.poll()));
            }
        }

        /**
         * Method to find the median value in a dataset
         * The median is the middle value in a sorted list of numbers
         * If the dataset has an even number of observations,
         * the median is the average of the two middle numbers
         *
         * @return the median value as a double
         */
        public double findMedian() {
            if (!maxHeap.isEmpty() && !minHeap.isEmpty() && minHeap.size() == maxHeap.size()) {
                return (double) (minHeap.peek() + maxHeap.peek()) / (double) 2;
            } else {
                return (double) maxHeap.peek();
            }
        }
    }

}
